import io
import os
import webbrowser
from tkinter import filedialog, messagebox
from tkinter import scrolledtext

import cv2
import numpy as np
import ttkbootstrap as ttk
from PIL import Image
from PyPDF2 import PdfReader, PdfWriter
from ttkbootstrap.constants import *


class PDFCompressPage(ttk.Frame):
    def __init__(self, master):
        super().__init__(master)
        self.pdf_files = []
        self.setup_ui()

    def setup_ui(self):
        # 提示信息
        info_frame = ttk.Frame(self)
        info_frame.pack(fill=X, padx=10, pady=5)

        info_label = ttk.Label(
            info_frame,
            text="提示：本功能仅对PDF中的图片进行压缩，不会压缩文本内容。",
            foreground="orange"
        )
        info_label.pack(side=LEFT)

        # 创建上传按钮
        upload_frame = ttk.Frame(self)
        upload_frame.pack(fill=X, pady=10)

        self.upload_btn = ttk.Button(
            upload_frame,
            text="上传PDF",
            command=self.upload_pdfs,
            style="primary.TButton"
        )
        self.upload_btn.pack(side=LEFT, padx=5)

        # 压缩质量设置
        quality_frame = ttk.Frame(self)
        quality_frame.pack(fill=X, padx=10, pady=5)

        ttk.Label(quality_frame, text="压缩质量:").pack(side=LEFT, padx=5)
        self.quality_var = ttk.IntVar(value=50)
        self.quality_scale = ttk.Scale(
            quality_frame,
            from_=1,
            to=100,
            orient=HORIZONTAL,
            variable=self.quality_var,
            length=500
        )
        self.quality_scale.pack(side=LEFT, padx=5)
        self.quality_label = ttk.Label(quality_frame, text="50%")
        self.quality_label.pack(side=LEFT, padx=5)
        self.quality_scale.configure(command=self.update_quality_label)

        # 显示已上传文件列表
        self.file_list = scrolledtext.ScrolledText(
            self,
            width=50,
            height=10,
            wrap=ttk.WORD
        )
        self.file_list.pack(padx=5, pady=5, fill=BOTH, expand=YES)
        self.file_list.config(state=DISABLED)

        # 压缩按钮
        self.compress_btn = ttk.Button(
            self,
            text="开始压缩",
            command=self.compress_pdfs,
            style="success.TButton"
        )
        self.compress_btn.pack(pady=10)

        # 进度条
        self.progress = ttk.Progressbar(
            self,
            mode='determinate',
            length=300
        )
        self.progress.pack(pady=10)

        # 压缩结果信息
        self.result_frame = ttk.LabelFrame(self, text="压缩结果")
        self.result_frame.pack(fill=X, padx=10, pady=10)

        self.result_text = scrolledtext.ScrolledText(
            self.result_frame,
            width=50,
            height=10,
            wrap=ttk.WORD
        )
        self.result_text.pack(padx=5, pady=5, fill=BOTH, expand=YES)
        self.result_text.config(state=DISABLED)

        # 导出文件信息
        self.export_frame = ttk.LabelFrame(self, text="导出文件")
        self.export_frame.pack(fill=X, padx=10, pady=10)

        self.export_path_label = ttk.Label(self.export_frame, text="未导出文件")
        self.export_path_label.pack(side=LEFT, padx=5)

        self.preview_export_btn = ttk.Button(
            self.export_frame,
            text="预览",
            command=self.preview_export,
            style="info.TButton",
            state=DISABLED
        )
        self.preview_export_btn.pack(side=LEFT, padx=5)

    def update_quality_label(self, value):
        self.quality_label.config(text=f"{int(float(value))}%")

    def update_result_text(self, text):
        self.result_text.config(state=NORMAL)
        self.result_text.delete(1.0, ttk.END)
        self.result_text.insert(ttk.END, text)
        self.result_text.config(state=DISABLED)

    def compress_image_with_opencv(self, image_np_array, quality):
        """使用OpenCV压缩图片质量"""
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
        _, img_encoded = cv2.imencode('.jpg', image_np_array, encode_param)
        return img_encoded.tobytes()

    def replace_image_in_pdf(self, pdf_reader, page_num, quality):
        """替换PDF页面中的图片"""
        page = pdf_reader.pages[page_num]
        resources = page['/Resources']
        image_count = 0

        if '/XObject' in resources:
            x_objects = resources['/XObject'].get_object()
            for obj_name in x_objects:
                if x_objects[obj_name]['/Subtype'] == '/Image':
                    try:
                        # 获取图像数据
                        img_stream = x_objects[obj_name]._data
                        pil_image = Image.open(io.BytesIO(img_stream))

                        # 转换为RGB模式
                        if pil_image.mode != 'RGB':
                            pil_image = pil_image.convert('RGB')

                        # 转换为numpy数组并调整颜色空间
                        image_np_array = np.array(pil_image)
                        image_np_array = cv2.cvtColor(image_np_array, cv2.COLOR_RGB2BGR)

                        # 压缩图片
                        compressed_image = self.compress_image_with_opencv(image_np_array, quality)
                        x_objects[obj_name]._data = compressed_image
                        image_count += 1

                    except Exception as e:
                        print(f"处理图片时出错: {str(e)}")
                        continue

        return page, image_count

    def upload_pdfs(self):
        file_paths = filedialog.askopenfilenames(
            filetypes=[("PDF文件", "*.pdf")]
        )
        if file_paths:
            self.pdf_files.extend(file_paths)
            self.update_file_list()
            self.update_result_text("已上传文件，请设置压缩质量后点击开始压缩。")

    def update_file_list(self):
        self.file_list.config(state=NORMAL)
        self.file_list.delete(1.0, ttk.END)

        for i, file_path in enumerate(self.pdf_files, 1):
            # 获取文件大小
            file_size = os.path.getsize(file_path)
            size_str = self.format_file_size(file_size)

            self.file_list.insert(ttk.END, f"{i}. {os.path.basename(file_path)} ({size_str})\n")

            # 添加操作按钮
            btn_frame = ttk.Frame(self.file_list)
            self.file_list.window_create(ttk.END, window=btn_frame)

            # 预览按钮
            preview_btn = ttk.Button(
                btn_frame,
                text="预览",
                command=lambda path=file_path: self.preview_pdf(path),
                style="info.TButton",
                width=5
            )
            preview_btn.pack(side=LEFT, padx=2)

            # 删除按钮
            delete_btn = ttk.Button(
                btn_frame,
                text="删除",
                command=lambda idx=i - 1: self.delete_pdf(idx),
                style="danger.TButton",
                width=5
            )
            delete_btn.pack(side=LEFT, padx=2)

            # 上移按钮
            if i > 1:
                up_btn = ttk.Button(
                    btn_frame,
                    text="↑",
                    command=lambda idx=i - 1: self.move_pdf(idx, -1),
                    style="primary.TButton",
                    width=2
                )
                up_btn.pack(side=LEFT, padx=2)

            # 下移按钮
            if i < len(self.pdf_files):
                down_btn = ttk.Button(
                    btn_frame,
                    text="↓",
                    command=lambda idx=i - 1: self.move_pdf(idx, 1),
                    style="primary.TButton",
                    width=2
                )
                down_btn.pack(side=LEFT, padx=2)

            self.file_list.insert(ttk.END, "\n")

        self.file_list.config(state=DISABLED)

    def format_file_size(self, size):
        """格式化文件大小显示"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024:
                return f"{size:.2f} {unit}"
            size /= 1024
        return f"{size:.2f} TB"

    def preview_pdf(self, file_path):
        try:
            webbrowser.open(file_path)
        except Exception as e:
            messagebox.showerror("错误", f"预览失败：{str(e)}")

    def delete_pdf(self, index):
        del self.pdf_files[index]
        self.update_file_list()
        self.update_result_text("已删除文件，请重新上传或开始压缩。")

    def move_pdf(self, index, direction):
        if 0 <= index + direction < len(self.pdf_files):
            self.pdf_files[index], self.pdf_files[index + direction] = \
                self.pdf_files[index + direction], self.pdf_files[index]
            self.update_file_list()

    def compress_pdfs(self):
        if not self.pdf_files:
            messagebox.showerror("错误", "请先上传PDF文件")
            return

        try:
            self.progress['value'] = 0
            self.update_idletasks()

            total_files = len(self.pdf_files)
            quality = self.quality_var.get()
            result_text = "压缩结果：\n\n"

            for i, pdf_path in enumerate(self.pdf_files):
                # 读取PDF
                reader = PdfReader(pdf_path)
                writer = PdfWriter()
                total_images = 0

                # 处理每一页
                for page_num in range(len(reader.pages)):
                    page, image_count = self.replace_image_in_pdf(reader, page_num, quality)
                    writer.add_page(page)
                    total_images += image_count

                # 设置输出路径
                base_name = os.path.splitext(os.path.basename(pdf_path))[0]
                default_name = f"{base_name}_compressed.pdf"

                output_path = filedialog.asksaveasfilename(
                    defaultextension=".pdf",
                    filetypes=[("PDF文件", "*.pdf")],
                    initialfile=default_name
                )

                if not output_path:
                    continue

                # 检查文件是否存在
                if os.path.exists(output_path):
                    if not messagebox.askyesno("确认", f"文件 {os.path.basename(output_path)} 已存在，是否覆盖？"):
                        continue

                # 保存压缩后的PDF
                with open(output_path, 'wb') as output_file:
                    writer.write(output_file)

                # 显示压缩结果
                original_size = os.path.getsize(pdf_path)
                compressed_size = os.path.getsize(output_path)
                compression_ratio = (1 - compressed_size / original_size) * 100

                result_text += (
                    f"文件: {os.path.basename(pdf_path)}\n"
                    f"原始大小: {self.format_file_size(original_size)}\n"
                    f"压缩后大小: {self.format_file_size(compressed_size)}\n"
                    f"压缩率: {compression_ratio:.2f}%\n"
                    f"处理的图片数量: {total_images}\n\n"
                )

                # 更新导出文件信息
                self.export_path_label.config(text=output_path)
                self.preview_export_btn.config(state=NORMAL)

                # 更新进度
                progress = (i + 1) / total_files * 100
                self.progress['value'] = progress
                self.update_idletasks()

            self.progress['value'] = 100
            self.update_idletasks()

            # 更新压缩结果
            self.update_result_text(result_text)

            messagebox.showinfo("成功", "PDF文件已压缩完成")

        except Exception as e:
            messagebox.showerror("错误", f"压缩失败：{str(e)}")
        finally:
            self.progress['value'] = 0

    def preview_export(self):
        try:
            webbrowser.open(self.export_path_label.cget("text"))
        except Exception as e:
            messagebox.showerror("错误", f"预览失败：{str(e)}")
